"""将原始形式的检测数据集格式转为COCO格式"""

import os
import json
import cv2

suffix = "train"
image_dir = f"/disk527/sdb1/a804_cbf/datasets/chang_e/{suffix}"
output_file = f"{suffix}.json"
images = []
annotations = []
for image_file in os.listdir(image_dir):
    if image_file.endswith(".png"):
        image_path = os.path.join(image_dir, image_file)
        image = cv2.imread(image_path)
        height, width = image.shape[:2]

        image_id = len(images) + 1
        image_info = {
            "id": image_id,
            "width": width,
            "height": height,
            "file_name": image_file,
        }
        images.append(image_info)

        with open(os.path.join(image_dir, image_file.replace("png", "txt")), "r") as f:
            lines = f.readlines()
            for line in lines:
                x1, y1, x2, y2 = map(float, line.strip().split(","))
                area = (x2 - x1) * (y2 - y1)

                annotation = {
                    "id": len(annotations) + 1,
                    "image_id": image_id,
                    "category_id": 0,
                    "bbox": [x1, y1, x2 - x1, y2 - y1],
                    "area": area,
                    "iscrowd": 0,
                }
                annotations.append(annotation)
coco_format = {
    "info": {"description": "luner crater dataset for chang'e 5 landing site"},
    "licenses": ["CC BY-NC-SA 4.0"],
    "images": images,
    "annotations": annotations,
    "categories": [{"id": 0, "name": "crater"}],
}


with open(output_file, "w") as f:
    json.dump(coco_format, f, indent=4)

print(f"COCO格式的标注文件已保存为 {output_file}")
